# ai_response.py
from sparkai.llm.llm import ChatSparkLLM, ChunkPrintHandler
from sparkai.core.messages import ChatMessage

SPARKAI_URL = 'wss://spark-api.xf-yun.com/v1.1/chat'
SPARKAI_APP_ID = 'ae6000ef'
SPARKAI_API_SECRET = 'OWQ5MjI3YjI3OTIyNjczNDg3YjAxZDFi'
SPARKAI_API_KEY = 'e263794f4d8fcf6b22f6f6f3f7de7390'
SPARKAI_DOMAIN = 'lite'

class AIResponse:
    def __init__(self):
        self.spark = ChatSparkLLM(
            spark_api_url=SPARKAI_URL,
            spark_app_id=SPARKAI_APP_ID,
            spark_api_key=SPARKAI_API_KEY,
            spark_api_secret=SPARKAI_API_SECRET,
            spark_llm_domain=SPARKAI_DOMAIN,
            streaming=False,
        )
    def get_response(self, message):
        messages = [ChatMessage(role="user", content=message)]
        handler = ChunkPrintHandler()
        response = self.spark.generate([messages], callbacks=[handler])
        response = self.extract_and_format_reply(response)
        return response
    
    def extract_and_format_reply(self, response):
        generations = response.generations
        for gen in generations:
            for sub_gen in gen:
                text_content = sub_gen.text
        
        return text_content

    def assist_write(self, message):
        messages = [ChatMessage(role="user", content='请根据以下内容生成一封邮件' + message)]
        handler = ChunkPrintHandler()
        response = self.spark.generate([messages], callbacks=[handler])
        response = self.extract_and_format_reply(response)
        return response
